import cv2
import numpy as np
import os
from imutils import paths
import pickle
import zipfile
import tempfile
import shutil
import glob
from tqdm import tqdm
import logging

# 确保日志配置正确
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
from PIL import UnidentifiedImageError


def get_size(file):
    """获取指定文件的大小（以MB为单位）"""
    size_bytes = os.path.getsize(file)
    return size_bytes / (1024 * 1024)


def createXY(train_folder, dest_folder, method='vgg', batch_size=64, max_images=None):
    """从训练文件夹或 zip 文件生成 X, y 特征与标签。

    特性：
    - 支持传入 .zip 文件或在当前工作目录自动查找 zip 并解压
    - 在遇到空或损坏的缓存时自动回退并删除缓存
    - 批次累积在循环内部，避免只写入最后一批的问题
    - 完全不依赖 VGG16 预训练权重，零下载
    """
    extracted_tmp = None

    def _cleanup():
        """清理临时解压目录"""
        try:
            if extracted_tmp and os.path.exists(extracted_tmp):
                shutil.rmtree(extracted_tmp)
                logging.info(f"清理临时解压目录: {extracted_tmp}")
        except Exception as e:
            logging.warning(f"清理临时目录失败: {e}")

    # 缓存文件路径
    x_file_path = os.path.join(dest_folder, f"X_{method}.pkl")
    y_file_path = os.path.join(dest_folder, f"y_{method}.pkl")
    legacy_x_path = os.path.join(dest_folder, "X.pkl")
    legacy_y_path = os.path.join(dest_folder, "y.pkl")

    # 处理zip文件解压
    try:
        if os.path.isfile(train_folder) and train_folder.lower().endswith('.zip'):
            logging.info(f"传入的 train_folder 是 zip 文件，解压到临时目录后使用: {train_folder}")
            extracted_tmp = tempfile.mkdtemp(prefix='createXY_')
            try:
                with zipfile.ZipFile(train_folder, 'r') as zf:
                    zf.extractall(extracted_tmp)
                train_folder = extracted_tmp
            except Exception as e:
                logging.exception(f"解压传入 zip 失败: {e}")
                _cleanup()
                return np.array([]), np.array([])
    except Exception as e:
        logging.warning(f"处理zip文件时出错: {e}")

    # 路径解析与验证
    if not os.path.exists(train_folder):
        alt = os.path.abspath(os.path.join(os.getcwd(), train_folder))
        if os.path.exists(alt):
            logging.info(f"train_folder 未直接找到，使用解析后路径: {alt}")
            train_folder = alt
        else:
            # 尝试在当前目录查找zip文件
            found_zip = False
            for z in glob.glob('*.zip'):
                tmp = tempfile.mkdtemp(prefix='createXY_')
                try:
                    with zipfile.ZipFile(z, 'r') as zf:
                        zf.extractall(tmp)
                    if any(paths.list_images(tmp)):
                        logging.info(f"在当前目录发现 zip 并解压可用数据: {z} -> {tmp}")
                        extracted_tmp = tmp
                        train_folder = tmp
                        found_zip = True
                        break
                    else:
                        shutil.rmtree(tmp)
                except Exception as e:
                    logging.warning(f"解压zip {z} 失败: {e}")
                    try:
                        shutil.rmtree(tmp)
                    except Exception:
                        pass
            if not found_zip:
                logging.error(f"train_folder 路径不存在，且未找到可用zip文件: {train_folder}")
                _cleanup()
                return np.array([]), np.array([])

    # 尝试读取缓存
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        try:
            x_size = os.path.getsize(x_file_path)
            y_size = os.path.getsize(y_file_path)
            logging.info(f"发现缓存（method={method}）：直接读取 {os.path.basename(x_file_path)}, {os.path.basename(y_file_path)}")
            logging.info(f"X文件大小:{get_size(x_file_path):.2f}MB")
            logging.info(f"y文件大小:{get_size(y_file_path):.2f}MB")
            
            if x_size == 0 or y_size == 0:
                logging.warning("检测到空缓存文件，忽略并删除")
                os.remove(x_file_path)
                os.remove(y_file_path)
            else:
                with open(x_file_path, 'rb') as fx:
                    X = pickle.load(fx)
                with open(y_file_path, 'rb') as fy:
                    y = pickle.load(fy)
                
                if np.array(X).size > 0 and np.array(y).size > 0:
                    _cleanup()
                    return X, y
                else:
                    logging.warning("缓存文件包含空数组，删除缓存")
                    os.remove(x_file_path)
                    os.remove(y_file_path)
        except Exception as e:
            logging.exception(f"读取缓存失败，忽略缓存并重新生成: {e}")
    elif os.path.exists(legacy_x_path) and os.path.exists(legacy_y_path):
        try:
            lx_size = os.path.getsize(legacy_x_path)
            ly_size = os.path.getsize(legacy_y_path)
            logging.info(f"发现旧版缓存：直接读取 {os.path.basename(legacy_x_path)}, {os.path.basename(legacy_y_path)}")
            
            if lx_size == 0 or ly_size == 0:
                logging.warning("检测到空旧版缓存文件，忽略并删除")
                os.remove(legacy_x_path)
                os.remove(legacy_y_path)
            else:
                with open(legacy_x_path, 'rb') as fx:
                    X = pickle.load(fx)
                with open(legacy_y_path, 'rb') as fy:
                    y = pickle.load(fy)
                
                if np.array(X).size > 0 and np.array(y).size > 0:
                    _cleanup()
                    return X, y
                else:
                    logging.warning("旧版缓存包含空数组，删除缓存")
                    os.remove(legacy_x_path)
                    os.remove(legacy_y_path)
        except Exception as e:
            logging.exception(f"读取旧版缓存失败，忽略缓存并重新生成: {e}")

    logging.info(f"读取所有图像，生成X和y（method={method}）")

    # 再次验证路径
    if not os.path.exists(train_folder):
        alt2 = os.path.abspath(os.path.join(os.getcwd(), os.path.basename(train_folder)))
        if os.path.exists(alt2):
            logging.info(f"train_folder 未直接找到，使用候选路径: {alt2}")
            train_folder = alt2
        else:
            logging.error(f"最终验证失败，路径不存在: {train_folder}")
            _cleanup()
            return np.array([]), np.array([])

    # 获取图像路径列表
    image_paths = list(paths.list_images(train_folder))
    if max_images and isinstance(max_images, int) and max_images > 0:
        image_paths = image_paths[:max_images]

    if len(image_paths) == 0:
        logging.error(f"在路径 {train_folder} 中未找到任何图片文件。请检查路径或数据组织（例如 data/train/dog/*.jpg）。")
        _cleanup()
        return np.array([]), np.array([])

    X = []
    y = []
    model = None

    # 构建特征提取模型
    if method == 'vgg':
        try:
            from tensorflow.keras.models import Model
            from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, GlobalMaxPooling2D
            from tensorflow.keras.utils import load_img, img_to_array
            
            # 手动构建VGG16结构（无预训练权重）
            input_shape = (224, 224, 3)
            inputs = Input(shape=input_shape)
            
            # 卷积块1
            x = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)
            x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
            x = MaxPooling2D((2, 2), strides=(2, 2))(x)
            
            # 卷积块2
            x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
            x = MaxPooling2D((2, 2), strides=(2, 2))(x)
            
            # 卷积块3
            x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
            x = MaxPooling2D((2, 2), strides=(2, 2))(x)
            
            # 卷积块4
            x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
            x = MaxPooling2D((2, 2), strides=(2, 2))(x)
            
            # 卷积块5
            x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
            x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
            x = MaxPooling2D((2, 2), strides=(2, 2))(x)
            
            # 全局最大池化
            outputs = GlobalMaxPooling2D()(x)
            model = Model(inputs=inputs, outputs=outputs)
            logging.info("VGG16特征提取器构建完成（无预训练权重，零下载）")
        except Exception as e:
            logging.exception(f"构建VGG16模型失败，将使用flat方法: {e}")
            method = 'flat'  # 降级为flat方法

    # 批次处理图像
    num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)
    for idx in tqdm(range(num_batches), desc="读取图像并提取特征"):
        batch_images = []
        batch_labels = []
        start = idx * batch_size
        end = min((idx + 1) * batch_size, len(image_paths))
        
        for i in range(start, end):
            image_path = image_paths[i]
            try:
                if method == 'vgg' and model is not None:
                    # VGG特征提取
                    img = load_img(image_path, target_size=(224, 224))
                    img = img_to_array(img)
                    img = img[:, :, ::-1]  # BGR转RGB
                    img -= [103.939, 116.779, 123.68]  # 标准化
                    batch_images.append(img)
                else:
                    # flat特征提取（灰度+展平）
                    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
                    if img is None:
                        raise Exception("无法读取图像")
                    img = cv2.resize(img, (32, 32))
                    batch_images.append(img)

                # 提取标签
                parent_folder = os.path.basename(os.path.dirname(image_path))
                train_root = os.path.basename(os.path.normpath(train_folder))
                label = parent_folder if parent_folder != train_root else os.path.basename(image_path).split('.')[0]
                batch_labels.append(label)
            except (UnidentifiedImageError, OSError, ValueError, Exception) as e:
                logging.warning(f"跳过无效图像: {image_path} -> {str(e)}")
                continue

        if len(batch_images) == 0:
            logging.info(f"批次 {idx} 无有效图像，跳过")
            continue

        # 特征提取
        batch_images = np.array(batch_images)
        if method == 'vgg' and model is not None:
            batch_features = model.predict(batch_images, verbose=0)
        else:
            batch_features = batch_images.reshape((batch_images.shape[0], -1))

        # 累积结果
        X.extend(batch_features)
        y.extend(batch_labels)

    # 处理结果
    X = np.array(X)
    y = np.array(y)

    # 标签转换为数字
    if y.size > 0 and (y.dtype == object or str(y.dtype).startswith('U')):
        unique_labels = sorted(list(set(y)))
        label_map = {label: idx for idx, label in enumerate(unique_labels)}
        logging.info(f"标签映射: {label_map}")
        y = np.array([label_map[label] for label in y], dtype=np.int64)

    logging.info(f"特征提取完成：X.shape={X.shape}, y.shape={y.shape}")

    # 保存缓存
    if X.size > 0 and y.size > 0:
        try:
            os.makedirs(dest_folder, exist_ok=True)  # 确保目标文件夹存在
            with open(x_file_path, 'wb') as f:
                pickle.dump(X, f)
            with open(y_file_path, 'wb') as f:
                pickle.dump(y, f)
            logging.info(f"缓存已保存到: {x_file_path} 和 {y_file_path}")
        except Exception as e:
            logging.exception(f"保存缓存失败: {e}")
    else:
        logging.warning("无有效特征数据，跳过缓存保存")

    _cleanup()
    return X, y